package cn._51doit.day04.partitioning;

import org.apache.flink.api.common.functions.RichMapFunction;
import org.apache.flink.api.java.functions.KeySelector;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.datastream.KeyedStream;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.sink.RichSinkFunction;

/**
 * @create: 2021-10-19 16:29
 * @author: 今晚打脑斧
 * @program: KeybyDemo
 * @Description:
 **/
public class KeybyDemo {
    public static void main(String[] args) throws Exception {
        Configuration configuration = new Configuration();
        configuration.setInteger("rest.port",8081);
        StreamExecutionEnvironment env = StreamExecutionEnvironment.createLocalEnvironmentWithWebUI(configuration);

        DataStreamSource<String> lines = env.socketTextStream("doit01", 8888);
        //用RichMapFunction是为了获得Subtask的位置,因为MapFunction没有这个方法
        SingleOutputStreamOperator<String> mappedStream = lines.map(new RichMapFunction<String, String>() {
            @Override
            public String map(String vallue) throws Exception {
                int indexOfThisSubtask = getRuntimeContext().getIndexOfThisSubtask();
                return vallue + " -> " + indexOfThisSubtask;
            }
        })/*.setParallelism(2)*/; //设置map的并行度为2 那么indexOfThisSubtask是0或者1

        //使用keyby的方式进行分区
        KeyedStream<String, String> keyedStream = mappedStream.keyBy(new KeySelector<String, String>() {
            @Override
            public String getKey(String value) throws Exception {
                return value.split(" ")[0];
            }
        });

        //将数据进行输出
        keyedStream.addSink(new RichSinkFunction<String>() {
            @Override
            public void invoke(String value, Context context) throws Exception {
                int indexOfThisSubtask = getRuntimeContext().getIndexOfThisSubtask();
                System.out.println(value+" -> "+indexOfThisSubtask);
            }
        })/*.setParallelism(3)*/;//设置sink的并行度为3 那么indexOfThisSubtask是0或者1或者2

        env.execute();

    }
}
